package middleware

import (
	"context"
	"net"
	"strings"

	"gitee.com/war2001/NiuMaEasyGo/global"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

// IPFilterConfig IP过滤配置
type IPFilterConfig struct {
	// 黑名单IP列表
	BlackList []string
	// 白名单IP列表
	WhiteList []string
	// 是否使用Redis存储黑名单
	UseRedis bool
	// Redis黑名单键
	BlackListKey string
	// Redis白名单键
	WhiteListKey string
}

// DefaultIPFilterConfig 默认IP过滤配置
var DefaultIPFilterConfig = IPFilterConfig{
	BlackList:    []string{},
	WhiteList:    []string{},
	UseRedis:     false,
	BlackListKey: "ip_blacklist",
	WhiteListKey: "ip_whitelist",
}

// IPFilter IP过滤中间件
func IPFilter(config ...IPFilterConfig) gin.HandlerFunc {
	// 使用默认配置或传入的配置
	ipFilterConfig := DefaultIPFilterConfig
	if len(config) > 0 {
		ipFilterConfig = config[0]
	}

	return func(c *gin.Context) {
		// 获取客户端IP
		clientIP := c.ClientIP()

		// 检查IP是否在白名单中
		if isIPInWhiteList(clientIP, ipFilterConfig) {
			c.Next()
			return
		}

		// 检查IP是否在黑名单中
		if isIPInBlackList(clientIP, ipFilterConfig) {
			zap.L().Warn("IP被拒绝访问",
				zap.String("ip", clientIP),
				zap.String("path", c.Request.URL.Path),
				zap.String("method", c.Request.Method),
			)
			ResponseForbidden(c, "您的IP地址被拒绝访问")
			c.Abort()
			return
		}

		c.Next()
	}
}

// isIPInBlackList 检查IP是否在黑名单中
func isIPInBlackList(ip string, config IPFilterConfig) bool {
	// 如果使用Redis存储黑名单
	if config.UseRedis && global.Redis != nil {
		// 检查IP是否在Redis黑名单中
		ctx := context.Background()
		exists, err := global.Redis.SIsMember(ctx, config.BlackListKey, ip).Result()
		if err != nil {
			zap.L().Error("检查IP黑名单失败", zap.Error(err))
			return false
		}
		if exists {
			return true
		}

		// 检查IP网段是否在Redis黑名单中
		members, err := global.Redis.SMembers(ctx, config.BlackListKey).Result()
		if err != nil {
			zap.L().Error("获取IP黑名单失败", zap.Error(err))
			return false
		}

		for _, cidr := range members {
			if strings.Contains(cidr, "/") && isIPInCIDR(ip, cidr) {
				return true
			}
		}

		return false
	}

	// 检查IP是否在内存黑名单中
	for _, blackIP := range config.BlackList {
		if strings.Contains(blackIP, "/") {
			// 如果是CIDR格式，检查IP是否在网段中
			if isIPInCIDR(ip, blackIP) {
				return true
			}
		} else {
			// 如果是单个IP，直接比较
			if ip == blackIP {
				return true
			}
		}
	}

	return false
}

// isIPInWhiteList 检查IP是否在白名单中
func isIPInWhiteList(ip string, config IPFilterConfig) bool {
	// 如果白名单为空，则不进行白名单检查
	if len(config.WhiteList) == 0 && (!config.UseRedis || global.Redis == nil) {
		return false
	}

	// 如果使用Redis存储白名单
	if config.UseRedis && global.Redis != nil {
		// 检查IP是否在Redis白名单中
		ctx := context.Background()
		exists, err := global.Redis.SIsMember(ctx, config.WhiteListKey, ip).Result()
		if err != nil {
			zap.L().Error("检查IP白名单失败", zap.Error(err))
			return false
		}
		if exists {
			return true
		}

		// 检查IP网段是否在Redis白名单中
		members, err := global.Redis.SMembers(ctx, config.WhiteListKey).Result()
		if err != nil {
			zap.L().Error("获取IP白名单失败", zap.Error(err))
			return false
		}

		for _, cidr := range members {
			if strings.Contains(cidr, "/") && isIPInCIDR(ip, cidr) {
				return true
			}
		}

		return false
	}

	// 检查IP是否在内存白名单中
	for _, whiteIP := range config.WhiteList {
		if strings.Contains(whiteIP, "/") {
			// 如果是CIDR格式，检查IP是否在网段中
			if isIPInCIDR(ip, whiteIP) {
				return true
			}
		} else {
			// 如果是单个IP，直接比较
			if ip == whiteIP {
				return true
			}
		}
	}

	return false
}

// isIPInCIDR 检查IP是否在CIDR网段中
func isIPInCIDR(ip, cidr string) bool {
	_, ipNet, err := net.ParseCIDR(cidr)
	if err != nil {
		zap.L().Error("解析CIDR失败", zap.String("cidr", cidr), zap.Error(err))
		return false
	}

	ipAddr := net.ParseIP(ip)
	if ipAddr == nil {
		zap.L().Error("解析IP失败", zap.String("ip", ip))
		return false
	}

	return ipNet.Contains(ipAddr)
}

// BlacklistIP 将IP加入黑名单
func BlacklistIP(ip string) error {
	// 如果Redis不可用，则跳过
	if global.Redis == nil {
		return nil
	}

	// 将IP加入黑名单
	ctx := context.Background()
	err := global.Redis.SAdd(ctx, DefaultIPFilterConfig.BlackListKey, ip).Err()
	if err != nil {
		zap.L().Error("将IP加入黑名单失败", zap.String("ip", ip), zap.Error(err))
		return err
	}

	zap.L().Info("IP已加入黑名单", zap.String("ip", ip))
	return nil
}

// WhitelistIP 将IP加入白名单
func WhitelistIP(ip string) error {
	// 如果Redis不可用，则跳过
	if global.Redis == nil {
		return nil
	}

	// 将IP加入白名单
	ctx := context.Background()
	err := global.Redis.SAdd(ctx, DefaultIPFilterConfig.WhiteListKey, ip).Err()
	if err != nil {
		zap.L().Error("将IP加入白名单失败", zap.String("ip", ip), zap.Error(err))
		return err
	}

	zap.L().Info("IP已加入白名单", zap.String("ip", ip))
	return nil
}

// RemoveIPFromBlacklist 从黑名单中移除IP
func RemoveIPFromBlacklist(ip string) error {
	// 如果Redis不可用，则跳过
	if global.Redis == nil {
		return nil
	}

	// 从黑名单中移除IP
	ctx := context.Background()
	err := global.Redis.SRem(ctx, DefaultIPFilterConfig.BlackListKey, ip).Err()
	if err != nil {
		zap.L().Error("从黑名单中移除IP失败", zap.String("ip", ip), zap.Error(err))
		return err
	}

	zap.L().Info("IP已从黑名单中移除", zap.String("ip", ip))
	return nil
}

// RemoveIPFromWhitelist 从白名单中移除IP
func RemoveIPFromWhitelist(ip string) error {
	// 如果Redis不可用，则跳过
	if global.Redis == nil {
		return nil
	}

	// 从白名单中移除IP
	ctx := context.Background()
	err := global.Redis.SRem(ctx, DefaultIPFilterConfig.WhiteListKey, ip).Err()
	if err != nil {
		zap.L().Error("从白名单中移除IP失败", zap.String("ip", ip), zap.Error(err))
		return err
	}

	zap.L().Info("IP已从白名单中移除", zap.String("ip", ip))
	return nil
}
